# user_behavior_markov_clustering.py
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
from sklearn.preprocessing import StandardScaler
from tqdm import tqdm  # 导入tqdm库

# 定义行为类型常量
BEHAVIOR_TYPES = ['pv', 'fav', 'cart', 'buy']


# 1. 数据准备
def load_data():
    print("1. 正在加载数据...")
    # 加载数据
    data = pd.read_parquet("../static/data/user_behavior_20141118_20141218_edit.parquet")

    # 转换时间格式
    data['time'] = pd.to_datetime(data['time'])
    data['hour'] = data['time'].dt.hour
    data['date'] = data['time'].dt.date

    print("✓ 数据加载完成")
    return data


# 2. 构建行为序列
def create_behavior_sequences(data):
    print("\n2. 正在构建行为序列...")
    # 按用户和时间排序
    data = data.sort_values(['user_id', 'time'])

    # 创建用户行为序列
    sequences = data.groupby('user_id')['behavior_type'].apply(list).reset_index()
    sequences.columns = ['user_id', 'behavior_sequence']

    # 计算序列长度
    sequences['sequence_length'] = sequences['behavior_sequence'].apply(len)

    # 过滤掉太短的序列
    sequences = sequences[sequences['sequence_length'] >= 3]

    print(f"✓ 行为序列构建完成，共 {len(sequences)} 个有效用户序列")
    return sequences


# 3. 构建转移矩阵
def build_transition_matrix(sequences):
    print("\n3. 正在构建转移矩阵...")
    # 初始化转移计数矩阵
    transition_counts = pd.DataFrame(0, index=BEHAVIOR_TYPES, columns=BEHAVIOR_TYPES)

    # 统计所有用户的转移次数
    for seq in tqdm(sequences['behavior_sequence'], desc="处理用户序列"):
        for i in range(len(seq) - 1):
            current = seq[i]
            next_ = seq[i + 1]
            transition_counts.loc[current, next_] += 1

    # 转换为转移概率矩阵
    transition_matrix = transition_counts.div(transition_counts.sum(axis=1), axis=1)
    transition_matrix = transition_matrix.fillna(0)

    print("✓ 转移矩阵构建完成")
    return transition_matrix


# 4. 特征工程
def extract_sequence_features(sequences, transition_matrix):
    print("\n4. 正在提取序列特征...")
    # 初始化特征列表
    features = []

    # 为每个用户序列提取特征
    for seq in tqdm(sequences['behavior_sequence'], desc="提取特征"):
        # 基本统计特征
        seq_length = len(seq)
        pv_count = seq.count('pv')
        fav_count = seq.count('fav')
        cart_count = seq.count('cart')
        buy_count = seq.count('buy')

        # 转移概率特征
        transition_probs = []
        for i in range(len(seq) - 1):
            current = seq[i]
            next_ = seq[i + 1]
            transition_probs.append(transition_matrix.loc[current, next_])

        avg_transition_prob = np.mean(transition_probs) if transition_probs else 0
        min_transition_prob = np.min(transition_probs) if transition_probs else 0
        max_transition_prob = np.max(transition_probs) if transition_probs else 0

        # 购买转化特征
        has_buy = 1 if 'buy' in seq else 0
        buy_position = seq.index('buy') / seq_length if has_buy else 0

        # 添加到特征列表
        features.append([
            seq_length,
            pv_count,
            fav_count,
            cart_count,
            buy_count,
            avg_transition_prob,
            min_transition_prob,
            max_transition_prob,
            has_buy,
            buy_position
        ])

    # 创建特征DataFrame
    feature_columns = [
        'seq_length', 'pv_count', 'fav_count', 'cart_count', 'buy_count',
        'avg_trans_prob', 'min_trans_prob', 'max_trans_prob',
        'has_buy', 'buy_position'
    ]
    features_df = pd.DataFrame(features, columns=feature_columns)

    # 标准化特征
    print("正在标准化特征...")
    scaler = StandardScaler()
    scaled_features = scaler.fit_transform(features_df)
    features_df = pd.DataFrame(scaled_features, columns=feature_columns)

    print("✓ 特征提取完成")
    return features_df


# 5. 马尔可夫聚类分析
def markov_clustering_analysis(features_df):
    print("\n5. 正在进行聚类分析...")
    # 使用轮廓系数确定最佳聚类数
    silhouette_scores = []
    ks = range(2, 8)

    print("正在计算不同聚类数的轮廓系数...")
    for k in tqdm(ks, desc="测试聚类数"):
        kmeans = KMeans(n_clusters=k, random_state=42)
        labels = kmeans.fit_predict(features_df)
        score = silhouette_score(features_df, labels)
        silhouette_scores.append(score)

    # 绘制轮廓系数图
    plt.figure(figsize=(10, 6))
    plt.plot(ks, silhouette_scores, 'bo-')
    plt.xlabel('Number of clusters')
    plt.ylabel('Silhouette Score')
    plt.title('Silhouette Score for KMeans Clustering')
    plt.grid()
    plt.savefig("../static/images/markov_silhouette_score.png")
    plt.close()  # 使用plt.close()代替plt.show()以避免阻塞

    # 选择最佳聚类数
    best_k = ks[np.argmax(silhouette_scores)]
    print(f"✓ 最佳聚类数确定: {best_k}")

    # 最终聚类
    print("正在进行最终聚类...")
    kmeans = KMeans(n_clusters=best_k, random_state=42)
    cluster_labels = kmeans.fit_predict(features_df)

    print("✓ 聚类分析完成")
    return cluster_labels, best_k


# 6. 可视化分析
def visualize_clusters(sequences, cluster_labels, best_k, transition_matrix):
    print("\n6. 正在生成可视化结果...")
    # 添加聚类标签到序列数据
    sequences['cluster'] = cluster_labels

    print("正在生成聚类转移矩阵图...")
    # 绘制每个聚类的行为转移图
    for cluster in tqdm(range(best_k), desc="生成聚类转移图"):
        cluster_sequences = sequences[sequences['cluster'] == cluster]['behavior_sequence']

        # 创建该聚类的转移矩阵
        cluster_transition = pd.DataFrame(0, index=BEHAVIOR_TYPES, columns=BEHAVIOR_TYPES)
        for seq in cluster_sequences:
            for i in range(len(seq) - 1):
                current = seq[i]
                next_ = seq[i + 1]
                cluster_transition.loc[current, next_] += 1

        # 转换为概率
        cluster_transition = cluster_transition.div(cluster_transition.sum(axis=1), axis=1).fillna(0)

        # 绘制转移图
        plt.figure(figsize=(8, 6))
        plt.title(f'Behavior Transition - Cluster {cluster}')
        sns.heatmap(cluster_transition, annot=True, cmap='Blues', fmt='.2f')
        plt.savefig(f"../static/images/markov_cluster_{cluster}_transition.png")
        plt.close()

    print("正在生成雷达图...")
    # 绘制聚类特征雷达图
    features_df = extract_sequence_features(sequences, transition_matrix)  # 重新生成features_df
    features_df['cluster'] = cluster_labels
    cluster_means = features_df.groupby('cluster').mean()

    # 雷达图设置
    categories = list(cluster_means.columns)
    N = len(categories)

    angles = [n / float(N) * 2 * np.pi for n in range(N)]
    angles += angles[:1]

    plt.figure(figsize=(10, 10))
    ax = plt.subplot(111, polar=True)

    for cluster in range(best_k):
        values = cluster_means.loc[cluster].values.flatten().tolist()
        values += values[:1]
        ax.plot(angles, values, linewidth=2, linestyle='solid', label=f'Cluster {cluster}')
        ax.fill(angles, values, alpha=0.25)

    plt.xticks(angles[:-1], categories)
    plt.title('Cluster Characteristics Radar Chart')
    plt.legend(loc='upper right', bbox_to_anchor=(1.3, 1.1))
    plt.savefig("../static/images/markov_cluster_radar.png")
    plt.close()

    print("✓ 可视化完成")


# 7. 主函数
def main():
    print("=" * 50)
    print("开始用户行为马尔可夫聚类分析")
    print("=" * 50)

    # 加载数据
    data = load_data()

    # 创建行为序列
    sequences = create_behavior_sequences(data)

    # 构建全局转移矩阵
    transition_matrix = build_transition_matrix(sequences)

    # 提取特征
    features_df = extract_sequence_features(sequences, transition_matrix)

    # 聚类分析
    cluster_labels, best_k = markov_clustering_analysis(features_df)

    # 可视化
    visualize_clusters(sequences, cluster_labels, best_k, transition_matrix)

    # 保存结果
    print("\n7. 正在保存结果...")
    sequences.to_csv("../static/data/user_behavior_clusters.csv", index=False)
    print("✓ 结果已保存到 user_behavior_clusters.csv")

    print("\n" + "=" * 50)
    print("分析完成!")
    print("=" * 50)


if __name__ == "__main__":
    main()
